src/mcp_agent/workflows/llm/augmented_llm.py (1)

598-652: Also surface strict in tracing and provider contracts.

To make runs diagnosable and consistent across providers, record strict and, optionally, whether a response_format is present.

Add to annotate_span_with_request_params:

         if hasattr(request_params, "metadata") and request_params.metadata:
             record_attributes(span, request_params.metadata, "request_params.metadata")
+        if hasattr(request_params, "strict"):
+            span.set_attribute("request_params.strict", bool(request_params.strict))
+        # Convenience boolean: did caller request structured output?
+        span.set_attribute(
+            "request_params.has_response_format",
+            bool(getattr(request_params, "metadata", {}) and request_params.metadata.get("response_format")),
+        )

Would you like me to open a follow-up to propagate strict in the non-Azure providers’ generate_structured paths?

src/mcp_agent/config.py (2)

13-13: Avoid import-time hard dependency on httpx in config module.

from httpx import URL at import time will raise if httpx isn’t installed, even when Anthropic isn’t used.

Use TYPE_CHECKING to gate the import:

+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from httpx import URL  # type: ignore

Then annotate with a forward ref (see next comment).

---

209-209: Prefer forward-ref type and alias parity.

• Use a forward ref to avoid runtime import: str | "URL" | None.
• Optional: add validation_alias for parity with other fields.

-    base_url: str | URL | None = Field(default=None)
+    base_url: str | "URL" | None = Field(
+        default=None,
+        validation_alias=AliasChoices("base_url", "ANTHROPIC_BASE_URL", "anthropic__base_url"),
+    )

src/mcp_agent/workflows/llm/augmented_llm_azure.py (2)

480-503: Type/semantics mismatch: response argument is a ChatCompletions, not a ResponseMessage.

The method operates on .choices[0].message, which exists on ChatCompletions. Update the annotation (and docstring) to avoid confusion and future errors.

-    def _annotate_span_for_completion_response(
-        self, span: trace.Span, response: ResponseMessage, turn: int
-    ) -> None:
+    def _annotate_span_for_completion_response(
+        self, span: trace.Span, response: ChatCompletions, turn: int
+    ) -> None:
         """Annotate the span with the completion response as an event."""

Optionally set event_name from response.choices[0].message.role as you already do.

---

132-139: Nit: docstring mentions “Azure OpenAI 5” — likely meant “Azure OpenAI (e.g., GPT‑4o‑mini)”.

Update to avoid confusion.

-        The default implementation uses Azure OpenAI 5 as the LLM.
+        The default implementation uses Azure OpenAI (e.g., GPT-4o-mini) as the LLM.

src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (1)

475-491: Harden tool_use extraction and diagnostics.

If parsing fails, we drop error context. Capture the last tool_use payload and include model name in the exception. Also accept string/dict inputs defensively.

             # Extract tool_use input and validate
-            for block in final.content:
+            last_tool_input = None
+            for block in final.content:
                 if (
                     getattr(block, "type", None) == "tool_use"
                     and getattr(block, "name", "") == "return_structured_output"
                 ):
                     data = getattr(block, "input", None)
+                    last_tool_input = data
                     try:
                         if isinstance(data, str):
                             return response_model.model_validate(json.loads(data))
                         return response_model.model_validate(data)
                     except Exception:
                         # Fallthrough to error
                         break
 
-            raise ValueError(
-                "Failed to obtain structured output from Anthropic response"
-            )
+            raise ValueError(
+                f"Failed to obtain/validate structured output from Anthropic response (model={model_name}); last_tool_input={to_string(last_tool_input)}"
+            )

src/mcp_agent/workflows/llm/augmented_llm_google.py (1)

264-272: Optional: surface strict intent.

Gemini ignores “strict” today; if params.strict is set, append a short system hint like “Only output JSON per the response schema—no prose.” Helps guard against extra text.

         config = types.GenerateContentConfig(
             max_output_tokens=params.maxTokens,
             temperature=params.temperature,
             stop_sequences=params.stopSequences or [],
-            system_instruction=self.instruction or params.systemPrompt,
+            system_instruction=(self.instruction or params.systemPrompt)
+            + (" Only output JSON per the response schema—no prose." if params.strict else ""),
         )

src/mcp_agent/workflows/llm/augmented_llm_openai.py (1)

474-482: Optional: include schema title and ensure deterministic ordering.

Some OpenAI models are sensitive to schema churn. Sorting properties and setting a stable “name” helps cache hits and reduces drift.

-            schema = response_model.model_json_schema()
+            schema = response_model.model_json_schema()
+            if "properties" in schema and isinstance(schema["properties"], dict):
+                schema["properties"] = dict(sorted(schema["properties"].items()))

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/utils/config.py : Configuration values such as quality_threshold, max_refinement_attempts, consolidation_interval, and evaluator_model_provider must be loaded from mcp_agent.config.yaml.

Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.

src/mcp_agent/workflows/llm/augmented_llm_azure.py (1)

545-576: Deep-copy payload, use project logger, and don’t override temperature in 400 fallback.

• Shallow copy risks nested mutation by the SDK.
• Prefer get_logger for consistency.
• Changing temperature to 1 can materially alter behavior; keep original value (prior feedback echoed here).

-        payload = request.payload.copy()
+        # Protect nested structures from mutation
+        import copy
+        payload = copy.deepcopy(request.payload)
@@
-        except HttpResponseError as e:
-            logger = logging.getLogger(__name__)
+        except HttpResponseError as e:
+            logger = get_logger(__name__)
@@
-            # Create a new payload with fallback values for commonly problematic parameters
-            fallback_payload = {**payload, "max_tokens": None, "temperature": 1}
+            # Retry with conservative fallback (preserve caller's temperature)
+            fallback_payload = {**payload, "max_tokens": None}

Add the missing import near the top if not present:

+import copy

src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (2)

475-492: Narrow exception handling and log validation/parse errors.

Catching Exception hides root causes and makes CI debugging hard. Catch json.JSONDecodeError and pydantic.ValidationError explicitly and log them.

-        import json
+        import json
+        from pydantic import ValidationError
...
-                    except Exception:
-                        # Fallthrough to error
-                        break
+                    except json.JSONDecodeError as e:
+                        self.logger.error("Failed to parse JSON structured output", exc_info=e)
+                        break
+                    except ValidationError as e:
+                        self.logger.error("Structured output failed model validation", exc_info=e)
+                        break
+                    except Exception as e:
+                        self.logger.error("Unexpected error validating structured output", exc_info=e)
+                        break

---

462-474: Don’t instantiate/stream AsyncAnthropic directly here—reuse the executor path to preserve provider routing (bedrock/vertex), telemetry, and avoid SDK context-manager mismatches.

This bypasses AnthropicCompletionTasks.request_completion_task, so config-driven routing, tracing, and our workflow_task semantics are skipped. Also, some Anthropic SDK versions don’t support async with client.messages.stream(...) causing __aenter__ failures in CI. Use the existing executor like generate() does.

-            # Call Anthropic directly (one-turn streaming for consistency)
-            base_url = None
-            if self.context and self.context.config and self.context.config.anthropic:
-                base_url = self.context.config.anthropic.base_url
-                api_key = self.context.config.anthropic.api_key
-                client = AsyncAnthropic(api_key=api_key, base_url=base_url)
-            else:
-                client = AsyncAnthropic()
-
-            async with client:
-                async with client.messages.stream(**args) as stream:
-                    final = await stream.get_final_message()
+            # Call Anthropic via the workflow task to honor provider routing & tracing
+            request = RequestCompletionRequest(
+                config=self.context.config.anthropic,
+                payload=args,
+            )
+            self._annotate_span_for_completion_request(span, request, 0)
+            final: Message = await self.executor.execute(
+                AnthropicCompletionTasks.request_completion_task,
+                ensure_serializable(request),
+            )
+            self._annotate_span_for_completion_response(span, final, 0)

If you prefer to keep direct SDK usage, at minimum gate by provider (avoid AsyncAnthropic for bedrock/vertex) and verify the installed SDK supports the async CM on messages.stream().

tests/workflows/llm/test_augmented_llm_google.py (1)

176-180: Update docstring to reflect native Gemini structured outputs (not Instructor).

-        Tests structured output generation using Instructor.
+        Tests structured output generation using Gemini native structured outputs.

tests/workflows/llm/test_augmented_llm_openai.py (1)

179-179: Remove redundant local import json; top-level import already present.

-        import json
+        # use module-level json import

Also applies to: 610-610

src/mcp_agent/workflows/llm/augmented_llm_azure.py (1)

548-576: Optional: add capped backoff for transient errors (429/5xx).

Exponential backoff with jitter and a small retry budget will improve resilience without hammering the endpoint. I can draft a focused utility if helpful.

src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (1)

440-456: Thread RequestParams.strict into the JSON Schema (disallow extras when strict).

When params.strict is true, set additionalProperties: false on the root schema so the model is steered to emit only declared fields.

-            schema = response_model.model_json_schema()
+            schema = response_model.model_json_schema()
+            # Tighten schema if strict requested
+            if getattr(params, "strict", False) and isinstance(schema, dict):
+                schema.setdefault("additionalProperties", False)

tests/workflows/llm/test_augmented_llm_anthropic.py (2)

241-253: Streaming SDK patching tightly couples tests to the direct SDK path. Consider stubbing the executor instead.

If you adopt the executor refactor above, update tests to patch llm.executor.execute to return a Message with a tool_use block instead of patching AsyncAnthropic.

Example change:

-        with patch(
-            "mcp_agent.workflows.llm.augmented_llm_anthropic.AsyncAnthropic"
-        ) as MockAsyncAnthropic:
-            ...
-            result = await AnthropicAugmentedLLM.generate_structured(
-                mock_llm, "Test query", TestResponseModel
-            )
+        mock_llm.executor.execute = AsyncMock(return_value=mock_message)
+        result = await mock_llm.generate_structured("Test query", TestResponseModel)

---

255-257: Call the bound method on the instance for clarity.

Minor readability improvement; current call works but is unconventional.

-            result = await AnthropicAugmentedLLM.generate_structured(
-                mock_llm, "Test query", TestResponseModel
-            )
+            result = await mock_llm.generate_structured("Test query", TestResponseModel)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.

tests/workflows/llm/test_augmented_llm_google.py (4)

176-181: Update stale docstring (no longer uses Instructor).

The test now validates native JSON structured outputs; reflect that to avoid confusion.

-        Tests structured output generation using Instructor.
+        Tests structured output generation using Gemini native structured outputs (application/json).

---

186-202: Good shift to native JSON; also assert JSON-mode config to prevent regressions.

Add assertions that the request set response_mime_type and the schema as expected.

         result = await mock_llm.generate_structured("Test query", TestResponseModel)

         # Assertions
         assert isinstance(result, TestResponseModel)
         assert result.name == "Test"
         assert result.value == 42
+        # Verify JSON-mode configuration was used
+        req = mock_llm.executor.execute.call_args[0][1]
+        cfg = req.payload["config"]
+        assert cfg.response_mime_type == "application/json"
+        assert cfg.response_schema in (
+            TestResponseModel,
+            TestResponseModel.model_json_schema(),
+        )

---

186-190: Deduplicate inline import json; move to module scope.

Inline imports are repeated; centralize at top for clarity.

-        import json

Additionally add at the top of this file:

# At file top-level (near other imports)
import json

Also applies to: 769-774

---

769-781: Mirror JSON-mode assertions for the mixed-message test.

This ensures both code paths set the correct response config.

         result = await mock_llm.generate_structured(messages, TestResponseModel)
         assert isinstance(result, TestResponseModel)
         assert result.name == "MixedTypes"
         assert result.value == 123
+        # Verify JSON-mode configuration on the request
+        req = mock_llm.executor.execute.call_args[0][1]
+        cfg = req.payload["config"]
+        assert cfg.response_mime_type == "application/json"
+        assert cfg.response_schema in (
+            TestResponseModel,
+            TestResponseModel.model_json_schema(),
+        )

src/mcp_agent/mcp/mcp_aggregator.py (4)

352-359: Config check: OK, but warning level may be noisy

The fallback path is correct. Consider downgrading the log to info to avoid warning fatigue when running without per-server configs.

---

361-369: Ensure O(1) membership by normalizing allowed_tools to a set

If callers pass a list, membership is O(n) despite the comment. Normalize once.

Apply:

-                    allowed_tools = self.context.server_registry.get_server_config(
-                        server_name
-                    ).allowed_tools
+                    cfg = self.context.server_registry.get_server_config(server_name)
+                    allowed_tools = (
+                        set(cfg.allowed_tools)
+                        if getattr(cfg, "allowed_tools", None) is not None
+                        else None
+                    )

---

553-555: Typo in span name ("get_capabilitites")

Fix spelling to keep traces searchable.

-            f"{self.__class__.__name__}.get_capabilitites"
+            f"{self.__class__.__name__}.get_capabilities"

---

311-317: Duplicate loading during create()

aenter() already initializes and loads servers; calling load_servers() again is redundant work.

-                logger.debug("Loading servers...")
-                await instance.load_servers()
-
-                logger.debug("MCPAggregator created and initialized.")
+                logger.debug("MCPAggregator created and initialized via __aenter__().")

tests/mcp/test_mcp_aggregator.py (1)

919-961: Filtering happy-path coverage: LGTM; minor nit

Good coverage for allowed set behavior. Tiny nit: mock_fetch_capabilities returns (None, ...); returning (server_name, ...) would more closely mirror production, though not functionally required.

src/mcp_agent/cli/cloud/commands/logger/tail/main.py (5)

316-318: Set a finite connect timeout for streaming.

Keep read timeout unlimited for SSE, but avoid hanging indefinitely on connect.

-        async with httpx.AsyncClient(timeout=None) as client:
+        # No read timeout for SSE; finite connect timeout for resiliency
+        async with httpx.AsyncClient(timeout=httpx.Timeout(connect=30.0, read=None)) as client:

---

451-454: YAML streaming: emit a single document per log, not a list.

Printing a list per event is awkward for consumers. Emit a mapping.

-    elif format == "yaml":
-        cleaned_entry = _clean_log_entry(log_entry)
-        print(yaml.dump([cleaned_entry], default_flow_style=False))
+    elif format == "yaml":
+        cleaned_entry = _clean_log_entry(log_entry)
+        print(yaml.safe_dump(cleaned_entry, default_flow_style=False))

---

478-486: Make _parse_log_level robust to non-string inputs.

Avoid AttributeError if level isn’t a string.

-def _parse_log_level(level: str) -> str:
+def _parse_log_level(level: Any) -> str:
     """Parse log level from API format to clean display format."""
-    if level.startswith("LOG_LEVEL_"):
+    if isinstance(level, str) and level.startswith("LOG_LEVEL_"):
         clean_level = level.replace("LOG_LEVEL_", "")
         if clean_level == "UNSPECIFIED":
             return "UNKNOWN"
         return clean_level
-    return level.upper()
+    return str(level).upper()

---

72-76: Leverage Typer choices to drop manual validation.

Define choices at the option level so invalid values are rejected before runtime; then you can remove the manual checks.

Example:

from typing import Literal

def tail_logs(
    ...,
    order_by: Optional[Literal["timestamp", "severity"]] = typer.Option(None, "--order-by", help="..."),
    format: Literal["text", "json", "yaml"] = typer.Option("text", "--format", help="..."),
):
    ...

Then the manual validations for --order-by and --format can be removed.

Also applies to: 139-143

---

33-34: Clarify argument help text for identifiers.

Be explicit about expected forms to reduce confusion.

-        help="Server ID, URL, or app configuration ID to retrieve logs for"
+        help="App ID (app_*), App Configuration ID (apcnf_*), or server URL to retrieve logs for"

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro